from matplotlib import pyplot as plt
import numpy as np
import torch
from model.vae import VAE
def save_image(image, filename):
    # 确保图像在 CPU 上
    image = image.to('cpu')
    # 转换 Tensor 为 numpy 数组并调整维度为 (H, W, C)
    image = image.detach().numpy().transpose(1, 2, 0)
    # 如果图像数据在 [-1, 1] 范围内，将其转换到 [0, 1]
    if image.min() < 0:
        image = (image + 1) / 2
    # 裁剪数值以确保它们在 [0, 1] 范围内
    image = np.clip(image, 0, 1)
    # 保存图像
    plt.imsave(filename, image)

    
def generate(label=None,pth=None, device='cuda'if torch.cuda.is_available() else 'cpu'):
    net=VAE()
    net.to(device)
    if pth is not None:
        net.load_state_dict(torch.load(pth, map_location=device), strict=False)
    decoder=net.decoder
    # 随机生成正态分布 TODO: 添加随机噪声
    mean=torch.zeros(size=[1,16],device=device,dtype=torch.float32)
    var=torch.ones(size=[1,16],device=device,dtype=torch.float32)
    img=decoder(mean,var)
    print(img)
    save_image(img[0]*10,'gen.jpg')
    

def main():
    pth=r'net_params.pth'
    generate(label=None, pth=pth)
    
if __name__ == '__main__':
    main()